package ebpf_inet

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"log/slog"
	"net"
	"os"
	"os/signal"
	"syscall"

	"github.com/cilium/ebpf/link"
	"github.com/cilium/ebpf/perf"
	"github.com/cilium/ebpf/rlimit"
	"github.com/k8spacket/k8spacket/broker"
	ebpf_tools "github.com/k8spacket/k8spacket/ebpf/tools"
	"github.com/k8spacket/k8spacket/modules"
)

/*
bpf2go - eBPF bytecode compiler & go files generator
args:
-cc clang - select C compiler
-target native - means get target platform from golang env GOARCH
-type event - name of type in C ebpf program to generate Go declaration
-go-package -package for output go file
bpf - identity name of generating program
./bpf/inet.bpf.c - C language source file
*/
//go:generate go run github.com/cilium/ebpf/cmd/bpf2go -cc clang -target native -type event -go-package ebpf_inet bpf ./bpf/inet.bpf.c

type InetEbpf struct {
	Broker broker.IBroker
}

func (inetEbpf *InetEbpf) Init() {

	slog.Info("INIT inet")
	// Allow the current process to lock more memory than the default for eBPF resources. Default value is 64KB
	// https://prototype-kernel.readthedocs.io/en/latest/bpf/troubleshooting.html#memory-ulimits
	// requires on kernels < 5.11 to remove memlock (error: failed to set memlock rlimit: operation not permitted)
	if err := rlimit.RemoveMemlock(); err != nil {
		slog.Error("[inet] Remove memlock", "Error", err)
	}

	// Load pre-compiled programs and maps into the kernel.
	objs := bpfObjects{}
	if err := loadBpfObjects(&objs, nil); err != nil {
		slog.Error("[inet] Loading objects", "Error", err)
	}
	defer objs.Close()

	// attach the eBPF program to the tracepoint sock/inet_sock_set_state
	ln, err := link.Tracepoint("sock", "inet_sock_set_state", objs.bpfPrograms.InetSockSetState, nil)
	if err != nil {
		slog.Error("[inet] Cannot attach tracepoint", "Error", err)
	}
	defer ln.Close()

	// create new reader for perf events
	rd, err := perf.NewReader(objs.bpfMaps.Events, os.Getpagesize())
	if err != nil {
		slog.Error("[inet] Creating perf event reader", "Error", err)
	}
	defer rd.Close()

	go func() {
		// bpfEvent is generated by bpf2go and represents perf event type in eBPF program
		var event bpfEvent
		for {
			record, err := rd.Read()
			if err != nil {
				if errors.Is(err, perf.ErrClosed) {
					slog.Info("[inet] Received signal, exiting..")
					return
				}
				slog.Error("[inet] Reading from reader", "Error", err)
				continue
			}

			// Parse the perf event into a go bpfEvent structure.
			if err := binary.Read(bytes.NewBuffer(record.RawSample), binary.LittleEndian, &event); err != nil {
				slog.Error("[inet] Parsing perf event", "Error", err)
				continue
			}

			distribute(event, inetEbpf)
		}
	}()

	// graceful shutdown
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	<-ctx.Done()

	slog.Info("[inet] Closed gracefully")
}

func distribute(event bpfEvent, inet *InetEbpf) {
	tcpEvent := modules.TCPEvent{
		Client: modules.Address{
			Addr: intToIP4(event.Saddr),
			Port: event.Sport},
		Server: modules.Address{
			Addr: intToIP4(event.Daddr),
			Port: event.Dport},
		TxB:     event.TxB,
		RxB:     event.RxB,
		DeltaUs: event.DeltaUs / 1000,
		Closed:  event.Closed}
	ebpf_tools.EnrichAddress(&tcpEvent.Client)
	ebpf_tools.EnrichAddress(&tcpEvent.Server)

	inet.Broker.TCPEvent(tcpEvent)
}

func intToIP4(ipNum uint32) string {
	ip := make(net.IP, 4)
	binary.LittleEndian.PutUint32(ip, ipNum)
	return ip.String()
}
